Add LFM2 and LFM2.5 support and testing#6428
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| lfm2_2_5_template = { | ||
| # LFM2.5 renders every tool call of a turn into a single `<|tool_call_start|>[...]<|tool_call_end|>` block, as a | ||
| # comma-separated list of Python-style calls: `[get_weather(city='Paris', days=3), ping()]`. The `open_pattern` | ||
| # therefore matches either the opening bracket or the comma separating two calls, and `close_pattern` matches the | ||
| # closing paren plus, for the last call, the trailing `]<|tool_call_end|>`. | ||
| # | ||
| # Argument values are rendered by the template's `format_arg_value` macro: strings single-quoted, mappings as JSON, | ||
| # everything else via Jinja's `string` filter. `string_delims` turns the single-quoted strings back into JSON | ||
| # strings so they parse. Note that this round-trip is lossy for Python literals the `string` filter emits but JSON | ||
| # cannot express: `flag=False` parses back as the string `"False"`, not `False`. That's inherent to the upstream | ||
| # template, not something we can recover here. | ||
| # | ||
| # The thinking field is named `thinking` (not `reasoning_content` as in other families) to match what the LFM2.5 | ||
| # template itself reads back off the message. | ||
| "defaults": {"role": "assistant"}, | ||
| "start_anchor": "<|im_start|>assistant\n", | ||
| "fields": { | ||
| "thinking": { | ||
| "open": "<think>", | ||
| "close_pattern": r"</think>\s*", | ||
| "content": "text", | ||
| }, | ||
| "tool_calls": { | ||
| "open_pattern": r"(?:<\|tool_call_start\|>\[|,\s*)(?P<name>\w+)\(", | ||
| "close_pattern": r"\)(?:\]<\|tool_call_end\|>)?\s*", | ||
| "repeats": True, | ||
| "content": "xml-inline", | ||
| "content_args": { | ||
| "tag_pattern": r"(?P<key>\w+)=(?P<value>'[^']*'|\{.*?\}|[^,)]*)", | ||
| "value_parser": {"name": "json", "args": {"string_delims": [["'", "'"]], "allow_non_json": True}}, | ||
| }, | ||
| "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, | ||
| }, | ||
| "content": { | ||
| "close_pattern": r"<\|im_end\|>\s*", | ||
| "content": "text", | ||
| }, | ||
| }, | ||
| } |
kashif
left a comment
kashif
left a comment
There was a problem hiding this comment.
curious if the CI passes
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e29c92649
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| "repeats": True, | ||
| "content": "xml-inline", | ||
| "content_args": { | ||
| "tag_pattern": r"(?P<key>\w+)=(?P<value>'[^']*'|\{.*?\}|[^,)]*)", |
There was a problem hiding this comment.
When an LFM2.5 tool call has an array argument or a nested mapping, this regex stops the value at the first comma or first }. The companion template can render valid calls such as foo(items=[1, 2]) or foo(cfg={"a":{"b":1},"c":2}), but response parsing turns those into truncated strings like "[1" or partial JSON before the tool is executed. This silently corrupts tool arguments for GRPO/response parsing; the value parser needs to consume balanced list/object values rather than splitting on commas/braces.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, reproduced: foo(items=[1, 2]) parsed back as {"items": "[1"}, and foo(cfg={"a": {"b": 1}, "c": 2}) as truncated JSON. The value pattern had no branch for [, and its \{.*?\} branch was non-greedy. Fixed by matching balanced brackets (quote-aware, up to three nesting levels, the regex-fixed-depth cap, noted in the comment). Lists, nested mappings, and lists-of-objects now round-trip.
Not sure if it's the best solution, but it seems to work.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit c4b6d94. Configure here.
| chunked_grad = model_chunked.lm_head.weight.grad.clone() | ||
|
|
||
| torch.testing.assert_close(chunked_grad, ref_grad, atol=5e-2, rtol=5e-2) | ||
| torch.testing.assert_close(chunked_grad, ref_grad, atol=1e-3, rtol=1e-3) |
There was a problem hiding this comment.
Missing tied-embedding backward xfail
Medium Severity
This PR adds LFM2 models to _CHUNKED_LM_HEAD_MODEL_IDS and tightens test_backward, but the PR notes tied-embedding models should xfail at temperature != 1.0. No xfail or skip on tie_word_embeddings appears, so CI can fail for LFM2 and other tied tiny models at temperature=0.7.
Reviewed by Cursor Bugbot for commit c4b6d94. Configure here.
2 participants
LFM2 sees roughly 1000 runs per week as per our telemetry, but the architecture/template isn't covered by our test suite at all. This adds both, for the two generations that behave differently:
trl-internal-testing/tiny-Lfm2ForCausalLM(fromLiquidAI/LFM2-1.2B), landed in transformers 4.54.0, below our floor.trl-internal-testing/tiny-Lfm2ForCausalLM-2.5(fromLiquidAI/LFM2.5-230M), gated ontransformers>=5.0.0, which itsTokenizersBackendtokenizer requires.Both are wired into the usual parametrized lists (
test_data_utils,test_utils,test_sft_trainer,test_dpo_trainer,test_chat_template_utils).Chat templates
The two generations differ in a way worth stating plainly:
toolsand thetoolrole, but never readsmessage["tool_calls"]: an assistant tool call renders as an empty turn. So it goes in the does-not-support tool-calling list, and gets alfm2_training.jinjapatch since its template has no{% generation %}markers.tool_calls, as<|tool_call_start|>[name(key=value, ...)]<|tool_call_end|>. It needs no training patch (its template already ships generation markers and is prefix-preserving), but it does get anlfm2_2_5_templateresponse template so tool calls can be parsed back. Only the new-style template is provided; the legacyresponse_schemais on its way out and isn't worth adding for a new family.One known limitation, inherent to the upstream template: argument values are rendered with Jinja's
stringfilter, soflag=Falseparses back as the string"False". Strings, numbers, and nested JSON round-trip fine.Adding LFM2 to
_CHUNKED_LM_HEAD_MODEL_IDSsurfaced a pre-existing issue. For models that tie their embeddings, the chunked LM head'sgrad_hiddenreacheslm_head.weightthrough the embedding, where the chunked path (fp32 accumulation) and the reference (whose.float()backward casts to bf16) round differently. Both end up several percent from the true fp32 gradient.It is not LFM2-specific, the behaviour splits cleanly on
tie_word_embeddingsacross five architectures (~1e-6 fp32 relative error tied vs ~1e-11 untied). Gemma and Cohere are tied too, and stay inside tolerance only because their gradients are ~40x smaller.Follow-ups?
get_training_chat_template. Highest-value next step.Note
Low Risk
Changes are mostly new Jinja templates, response parsing, and test/param wiring for LFM2 families; no auth, data, or core training algorithm changes beyond template auto-selection.
Overview
Adds first-class support for Liquid LFM2 and LFM2.5 in TRL’s chat-template stack and test matrix.
LFM2 gets bundled
lfm2.jinja/lfm2_training.jinja: the training patch splits assistant turns so<|im_start|>assistant\nstays outside{% generation %}while content and<|im_end|>are masked for assistant-only SFT. It is registered for automatic training-template swap but classified as not supporting standardtool_calls(tools/tool roles render; assistanttool_callsare ignored).LFM2.5 adds
lfm2_2_5.jinjaplus a newlfm2_2_5_templateresponse parser (thinking + Python-style<|tool_call_start|>[fn(a=1)]<|tool_call_end|>blocks).add_response_schemamaps this family to the new-style template only (transformers ≥ 5.13) and errors on older versions without a legacy schema.Tiny checkpoint generators (
lfm2_for_causal_lm.py,lfm2_for_causal_lm_2_5.py) and parametrized tests wire both models into SFT/DPO, data utils, chat-template utilities, and chunked LM-head coverage (LFM2.5 gated on transformers ≥ 5.0).Docs and
trl/chat_templates/README.mddocument the two generations’ tool-calling behavior.Separately,
test_backwardfor chunked LM head now runs in float32 with tighter tolerances so tied-embedding bf16 noise does not mask gradient mismatches (issue surfaced when LFM2 joined the model list).Reviewed by Cursor Bugbot for commit f7626c4. Bugbot is set up for automated code reviews on this repo. Configure here.